import numpy as np
import matplotlib.pyplot as plt

# 加载数据
data_path = './dataset/channel_data/channel_data.npz'
print(f"加载数据: {data_path}")
data = np.load(data_path)

# 查看所有键
print(f"数据包含的键: {list(data.keys())}")

# 获取形状信息
channel_features = data['channel_features']  # [n_windows, window_size, n_clusters, 5]
probe_weights = data['probe_weights']        # [n_windows, window_size, n_probes]
window_indices = data['window_indices']      # 窗口索引

print(f"数据总体形状:")
print(f"- 信道特征: {channel_features.shape}")
print(f"- 探头权重: {probe_weights.shape}")
print(f"- 窗口索引: {window_indices.shape}")

# 查看单个窗口的结构
window_idx = 0  # 第一个窗口
window_data = channel_features[window_idx]  # [20, 25, 5]
print(f"\n单个窗口形状: {window_data.shape}")
print(f"- 窗口大小(时间点数): {window_data.shape[0]}")
print(f"- 每个时间点的簇数: {window_data.shape[1]}")
print(f"- 每个簇的特征数: {window_data.shape[2]}")

# 查看单个时间点的结构
time_idx = 0  # 第一个时间点
time_data = window_data[time_idx]  # [25, 5]
print(f"\n单个时间点形状: {time_data.shape}")

# 查看每个特征的含义
print(f"\n每个簇的5个特征含义:")
print(f"- 特征0: 水平到达角 (horizontal_aoa)")
print(f"- 特征1: 垂直到达角 (vertical_aoa)")
print(f"- 特征2: 接收功率 (received_power)")
print(f"- 特征3: 水平扩展角 (ASA)")
print(f"- 特征4: 垂直扩展角 (ZSA)")

# 打印第一个时间点的前5个簇数据
print(f"\n第一个时间点的前5个簇数据:")
for cluster_idx in range(5):
    cluster_data = time_data[cluster_idx]
    print(f"簇 {cluster_idx+1}:")
    print(f"  水平到达角: {cluster_data[0]:.2f}")
    print(f"  垂直到达角: {cluster_data[1]:.2f}")
    print(f"  接收功率: {cluster_data[2]:.2f}")
    print(f"  水平扩展角: {cluster_data[3]:.4f}")
    print(f"  垂直扩展角: {cluster_data[4]:.4f}")

# 查看不同特征的数值范围
print(f"\n特征的数值范围:")
for feature_idx in range(5):
    feature_data = channel_features[:, :, :, feature_idx].flatten()
    print(f"特征 {feature_idx}: min={feature_data.min():.2f}, max={feature_data.max():.2f}, mean={feature_data.mean():.2f}")

# 可视化第一个窗口的部分数据
plt.figure(figsize=(15, 10))

# 1. 可视化水平到达角
plt.subplot(2, 2, 1)
plt.imshow(window_data[:, :, 0], aspect='auto', cmap='viridis')
plt.colorbar(label='Horizontal AOA (degrees)')
plt.title('Horizontal Arrival Angle in Window (Timesteps x Clusters)')
plt.xlabel('Cluster Index')
plt.ylabel('Timestep Index')

# 2. 可视化垂直到达角
plt.subplot(2, 2, 2)
plt.imshow(window_data[:, :, 1], aspect='auto', cmap='viridis')
plt.colorbar(label='Vertical AOA (degrees)')
plt.title('Vertical Arrival Angle in Window (Timesteps x Clusters)')
plt.xlabel('Cluster Index')
plt.ylabel('Timestep Index')

# 3. 可视化接收功率
plt.subplot(2, 2, 3)
plt.imshow(window_data[:, :, 2], aspect='auto', cmap='inferno')
plt.colorbar(label='Received Power (dB)')
plt.title('Received Power in Window (Timesteps x Clusters)')
plt.xlabel('Cluster Index')
plt.ylabel('Timestep Index')

# 4. 可视化前5个簇在时间上的变化
plt.subplot(2, 2, 4)
for cluster_idx in range(5):
    plt.plot(window_data[:, cluster_idx, 2], label=f'Cluster {cluster_idx+1}')
plt.title('Received Power of Top 5 Clusters Over Time')
plt.xlabel('Timestep Index')
plt.ylabel('Received Power (dB)')
plt.legend()

plt.tight_layout()
plt.savefig('./dataset/channel_data/data_visualization.png')
print(f"\n可视化图像已保存到: ./dataset/channel_data/data_visualization.png") 